package cn.caplike.demo.kafka.producer.partition;

import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.serialization.StringSerializer;

import java.util.Map;
import java.util.Objects;
import java.util.Properties;

public class ProducerCustomPartitioner {

    public static void main(String[] args) {
        // 在控制台监听一个 topic: ./kafka-console-consumer.sh --bootstrap-server big-data-1:9092 big-data-2:9092 big-data-3:9092 --topic first

        // 配置
        final Properties properties = new Properties();

        // 连接集群 bootstrap.servers
        properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "big-data-1:9092,big-data-2:9092,big-data-3:9092");
        // 指定对应 key 和 value 的序列化类型 key.serializer
        properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        // 定义自定义分区器
        properties.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, MyPartitioner.class.getName());

        // 创建生产者对象
        try (final KafkaProducer<String, String/* "", hello */> producer = new KafkaProducer<>(properties)) {
            // 带回调的异步发送数据
            producer.send(new ProducerRecord<>("first","", "caplike"), new Callback() {
                @Override
                public void onCompletion(RecordMetadata metadata, Exception exception) {
                    if (Objects.isNull(exception)) {
                        System.out.println("topic: " + metadata.topic());
                        System.out.println("partition: " + metadata.partition());
                    } else {
                        throw new RuntimeException(exception);

                        //在 Kafka 的 Java 客户端中，如果将消息发送到一个不存在的分区，不会在回调函数中触发异常逻辑。回调函数只会在消息发送完成后被调用，而不会报告分区不存在的异常。
                        //实际上，Kafka 的生产者在发送消息时，并不会在客户端检查分区是否存在，而是将消息发送到 Kafka 服务器，由服务器进行处理。如果发送到一个不存在的分区，Kafka 服务器会返回一个错误响应，但这个错误响应不会被传递给客户端的回调函数。
                        //因此，要在生产者中检测分区是否存在，你可以使用 Kafka 的元数据 API 查询分区信息，然后根据查询结果来确定消息发送的目标分区是否存在。
                        //以下是一个示例代码片段，展示了如何使用 Kafka 的元数据 API 查询分区信息并处理不存在分区的情况：
                        //
                        //import org.apache.kafka.clients.admin.*;
                        //import org.apache.kafka.clients.producer.*;
                        //
                        //import java.util.*;
                        //import java.util.concurrent.ExecutionException;
                        //
                        //public class KafkaProducerExample {
                        //    public static void main(String[] args) throws ExecutionException, InterruptedException {
                        //        Properties props = new Properties();
                        //        // 设置 Kafka 服务器地址等配置信息
                        //
                        //        // 创建 Kafka 管理员客户端
                        //        AdminClient adminClient = AdminClient.create(props);
                        //
                        //        String topic = "your_topic";
                        //        String key = "your_key";
                        //        String value = "your_value";
                        //        int partition = 999; // 不存在的分区
                        //
                        //        // 查询主题的分区信息
                        //        DescribeTopicsResult topicsResult = adminClient.describeTopics(Collections.singleton(topic));
                        //        TopicDescription topicDescription = topicsResult.values().get(topic).get();
                        //        List<TopicPartitionInfo> partitions = topicDescription.partitions();
                        //
                        //        // 检查目标分区是否存在
                        //        boolean partitionExists = partitions.stream().anyMatch(p -> p.partition() == partition);
                        //        if (!partitionExists) {
                        //            System.err.println("目标分区不存在");
                        //            // 处理分区不存在的情况
                        //            return;
                        //        }
                        //
                        //        // 创建 Kafka 生产者
                        //        Producer<String, String> producer = new KafkaProducer<>(props);
                        //
                        //        ProducerRecord<String, String> record = new ProducerRecord<>(topic, partition, key, value);
                        //
                        //        // 发送消息，并设置回调函数
                        //        producer.send(record, new Callback() {
                        //            @Override
                        //            public void onCompletion(RecordMetadata metadata, Exception exception) {
                        //                if (exception != null) {
                        //                    System.err.println("消息发送失败：" + exception.getMessage());
                        //                } else {
                        //                    System.out.println("消息发送成功，分区：" + metadata.partition() + "，偏移量：" + metadata.offset());
                        //                }
                        //            }
                        //        });
                        //
                        //        producer.close();
                        //        adminClient.close();
                        //    }
                        //}
                        //在上述示例中，我们使用 Kafka 的元数据 API 查询主题的分区信息，然后检查目标分区是否存在。如果分区不存在，我们可以在相应的处理逻辑中进行处理。
                        //请注意，在查询分区信息时，需要使用 Kafka 管理员客户端（AdminClient），因此需要添加相应的依赖。
                    }
                }
            }); // 会自动创建 topic, 发送一条消息, 所有消费者都能收到
        }

    }

    public static class MyPartitioner implements Partitioner {

        @Override
        public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {

            // 获取数据
            final String msg = value.toString();

            int partition;
            if (msg.contains("like")) {
                partition = 999; // 不存在的分区
            } else {
                partition = 0;
            }

            return partition;
        }

        @Override
        public void close() {

        }

        @Override
        public void configure(Map<String, ?> configs) {

        }
    }
}


